import torch
import torchvision
import torch
import torchvision
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear, CrossEntropyLoss
from torch.utils.data import DataLoader

vgg16 = torchvision.models.vgg16(pretrained=False)
# 保存方式1 (保存是模型结构+模型参数) 加载的时候需要有定义的模型的类才能加载
# torch.save(vgg16, 'vgg16_mothod1.pth')

# 保存方式2 (模型参数）
# torch.save(vgg16.state_dict(), 'vgg16_method2.pth')

class module(nn.Module):

    def __init__(self):

        super(module, self).__init__()
        self.model1 = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 20)

        )

    def forward(self, x):

        x = self.model1(x)
        return x

model = module()
torch.save(model, 'test.pth')
